#!/usr/bin/env python3.8.13
# -*- coding: utf-8 -*-

# @FileName   : espnet_conformer_inference.py
# @Description：espnet 语音识别
import time
import os
import csv

import soundfile
# from asr_inference import Speech2Text
from espnet2.bin.asr_inference import Speech2Text
import logging
logging.basicConfig(level=logging.INFO)

class EspnetASR:
    """
    espnet预训练模型和配置文件设置, 识别
    """
    def __init__(self):
        dir_path = "/home/gyf/pkg/xxgg/github/ai_app/cmcc/docs/cmcc_2022/conformer_espnet/"
        dir_path = "/public/ai_platform/models/cmcc_2022/conformer_pytorch/"
        self._path_asr_config = dir_path+"exp/asr_train_asr_conformer3_raw_char_batch_bins4000000_accum_grad4_sp/config.yaml"
        self._path_asr_file = dir_path+"exp/asr_train_asr_conformer3_raw_char_batch_bins4000000_accum_grad4_sp/valid.acc.ave_10best.pth"

        self._speech2text = Speech2Text(asr_train_config=self._path_asr_config, asr_model_file=self._path_asr_file)

    def speech_to_text(self, wav_path):
        """
        音频文件识别函数
        :param wav_path: 需要进行识别的音频文件的路径
        :return result_recog: 返回识别的字符串
        """
        # self.select_asr_model()
        audio, rate = soundfile.read(wav_path)
        t1 = time.perf_counter()
        result_all = self._speech2text(audio)
        t2 = time.perf_counter()
        print(os.path.basename(wav_path), result_all[0][0], (t2-t1)*1000, "ms")
        result_text = result_all[0][0]
        return result_text

def main():
    # wav_path = "/home/gyf/pkg/xxgg/github/ai_app/cmcc/asr_conformer/speech_recognition/wav/BAC009S0764W0121.wav"
    # wave_path = "wav/dalang.wav"
    er = EspnetASR()
    print("##############")
    # print(er.speech_to_text(wav_path))

    dataset_dir = "/public/ai_platform/models/cmcc_2022/dataset_aishell"
    test_set = "dev" # "test": total time cost: 25498543.0164868
    save_csv = "result_dev_pth.csv"

    T1 = time.perf_counter()
    for root, dirs, files in os.walk(os.path.join(dataset_dir, test_set)):
        for file in files:
            if file.endswith(".wav") and file != "BAC009S0731W0379.wav" and file!="BAC009S0756W0343.wav": #  and file not in csv_dict.keys()
                # print(file)
                wav_file = os.path.join(root, file)
                T10 = time.perf_counter()
                result = er.speech_to_text(wav_file)
                T2 = time.perf_counter()
                # print(file, result, (T2-T10)*1000)
                with open(save_csv, "a+", encoding="utf-8-sig", newline="") as csv_file:
                    writer = csv.writer(csv_file)
                    writer.writerow((file, result, (T2-T1)*1000))
    Tn = time.perf_counter()
    print((Tn-T1)*1000)


if __name__ == "__main__":
    main()